import json
import time
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from wiki_search import WikiSearch

class WikiSearchAPI:
    def __init__(self, data_dir: str = "wiki_chunks"):
        """初始化Wiki搜索API
        
        Args:
            data_dir: 存放wiki分片数据的目录
        """
        self.searcher = WikiSearch(data_dir)
    
    def search(
        self,
        query: str,
        top_k: int = 5,
        min_score: float = 0.3
    ) -> Tuple[bool, str, List[Dict], str]:
        """执行Wiki内容搜索
        
        Args:
            query: 搜索查询文本
            top_k: 返回结果数量(1-20)
            min_score: 最低相关性分数阈值(0-1)
            
        Returns:
            Tuple[success, message, results, timestamp]
        """
        timestamp = datetime.utcnow().isoformat()
        
        try:
            # 参数验证
            if not query or not isinstance(query, str):
                return False, "查询文本不能为空", [], timestamp
            
            top_k = max(1, min(20, top_k))
            min_score = max(0, min(1, min_score))
            
            start_time = time.time()
            
            # 获取查询向量
            embedding = self.searcher.api.get_embedding(query)
            if not embedding:
                return False, "无法生成查询向量", [], timestamp
            
            # 执行搜索
            raw_results = self.searcher.db.search(embedding, top_k=top_k)
            
            # 过滤低于阈值的結果
            results = [
                r for r in raw_results 
                if r['score'] >= min_score
            ]
            
            elapsed = time.time() - start_time
            message = f"找到{len(results)}条结果 (耗时{elapsed:.2f}s)"
            
            return True, message, results, timestamp
            
        except Exception as e:
            return False, f"搜索失败: {str(e)}", [], timestamp
    
    def save_search_results(
        self,
        results: List[Dict],
        output_file: str = "search_results.json"
    ) -> bool:
        """保存搜索结果到JSON文件
        
        Args:
            results: 搜索结果列表
            output_file: 输出文件路径
            
        Returns:
            是否保存成功
        """
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(results, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"保存搜索结果失败: {str(e)}")
            return False


# 使用示例
if __name__ == "__main__":
    # 初始化API
    wiki_api = WikiSearchAPI()
    
    # 执行搜索
    success, message, results, timestamp = wiki_api.search(
        query="星穹列车成员",
        top_k=3,
        min_score=0.4
    )
    
    # 处理结果
    if success:
        print(f"搜索成功: {message}")
        print(f"时间戳: {timestamp}")
        print("搜索结果:")
        for i, result in enumerate(results, 1):
            print(f"\n结果 #{i}:")
            print(f"相关度: {result['score']:.3f}")
            print(f"内容: {result['text'][:200]}...")
            print(f"来源: {result['source_file']}")
            print(f"字符数: {result.get('char_length', 'N/A')}")
        
        # 保存结果
        # wiki_api.save_search_results(results)
    else:
        print(f"搜索失败: {message}")
